import torch
import torch.nn as nn


def not_freeze(name, freeze):
    for f in freeze:
        if f in name:
            return False
    return True

def compute_ewc_loss(ewc_pack, method = 'L2',  freeze = []):
    ewc_enable = ewc_pack['ewc_enable']
    if not ewc_enable:  # 如果没有启用 EWC，直接返回零损失
        return torch.zeros(1, device=ewc_pack['device'])
    
    fisher_matrix = ewc_pack['fisher_matrix']
    old_modelpara = ewc_pack['old_modelpara']
    model = ewc_pack['model']
    ewc_lambda = ewc_pack['ewc_lambda']
    ewc_loss = torch.zeros(1, device=ewc_pack['device'])
    #task_grads = ewc_pack['task_grads']
    
    device = ewc_pack['device']
    

    if method == 'L2':
        for name, param in model.named_parameters():
            # if ('cv3' in name or 'bn3' in name or 'm.3' in name): # and '24' not in name  :
            if not_freeze(name, freeze) and '24' not in name:
                ewc_loss += torch.sum(
                    fisher_matrix[name].detach().to(ewc_pack['device']) 
                    * ((param - old_modelpara[name].detach().to(ewc_pack['device'])) ** 2) 
                                                    )
        
        ewc_loss *= (ewc_lambda / 2) 
        return ewc_loss

    
        return ewc_loss * ewc_lambda




















